/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.core.construction;



import com.google.auto.service.AutoService;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.FunctionSpec;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.IsBounded;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.ReadPayload;

import com.bff.gaia.unified.model.pipeline.v1.RunnerApi.SdkFunctionSpec;

import com.bff.gaia.unified.sdk.io.BoundedSource;

import com.bff.gaia.unified.sdk.io.Read;

import com.bff.gaia.unified.sdk.io.Read.Unbounded;

import com.bff.gaia.unified.sdk.io.Source;

import com.bff.gaia.unified.sdk.io.UnboundedSource;

import com.bff.gaia.unified.sdk.runners.AppliedPTransform;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.util.SerializableUtils;

import com.bff.gaia.unified.sdk.values.PBegin;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.com.google.protobuf.ByteString;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.com.google.protobuf.InvalidProtocolBufferException;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableMap;



import java.io.IOException;

import java.util.Collections;

import java.util.Map;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;



/**

 * Methods for translating {@link Read.Bounded} and {@link Read.Unbounded} {@link PTransform

 * PTransformTranslation} into {@link ReadPayload} protos.

 */

public class ReadTranslation {

  private static final String JAVA_SERIALIZED_BOUNDED_SOURCE = "urn:unified:java:boundedsource:v1";

  private static final String JAVA_SERIALIZED_UNBOUNDED_SOURCE = "urn:unified:java:unboundedsource:v1";



  public static ReadPayload toProto(Read.Bounded<?> read, SdkComponents components) {

    return ReadPayload.newBuilder()

        .setIsBounded(IsBounded.Enum.BOUNDED)

        .setSource(toProto(read.getSource(), components))

        .build();

  }



  public static ReadPayload toProto(Unbounded<?> read, SdkComponents components) {

    return ReadPayload.newBuilder()

        .setIsBounded(IsBounded.Enum.UNBOUNDED)

        .setSource(toProto(read.getSource(), components))

        .build();

  }



  public static SdkFunctionSpec toProto(Source<?> source, SdkComponents components) {

    if (source instanceof BoundedSource) {

      return toProto((BoundedSource) source, components);

    } else if (source instanceof UnboundedSource) {

      return toProto((UnboundedSource<?, ?>) source, components);

    } else {

      throw new IllegalArgumentException(

          String.format("Unknown %s type %s", Source.class.getSimpleName(), source.getClass()));

    }

  }



  private static SdkFunctionSpec toProto(BoundedSource<?> source, SdkComponents components) {

    return SdkFunctionSpec.newBuilder()

        .setEnvironmentId(components.getOnlyEnvironmentId())

        .setSpec(

            FunctionSpec.newBuilder()

                .setUrn(JAVA_SERIALIZED_BOUNDED_SOURCE)

                .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(source)))

                .build())

        .build();

  }



  public static BoundedSource<?> boundedSourceFromProto(ReadPayload payload)

      throws InvalidProtocolBufferException {

    checkArgument(payload.getIsBounded().equals(IsBounded.Enum.BOUNDED));

    return (BoundedSource<?>)

        SerializableUtils.deserializeFromByteArray(

            payload.getSource().getSpec().getPayload().toByteArray(), "BoundedSource");

  }



  public static <T> BoundedSource<T> boundedSourceFromTransform(

      AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform)

      throws IOException {

    return (BoundedSource<T>) boundedSourceFromProto(getReadPayload(transform));

  }



  public static <T, CheckpointT extends UnboundedSource.CheckpointMark>

      UnboundedSource<T, CheckpointT> unboundedSourceFromTransform(

          AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform)

          throws IOException {

    return (UnboundedSource<T, CheckpointT>) unboundedSourceFromProto(getReadPayload(transform));

  }



  private static <T> ReadPayload getReadPayload(

      AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform)

      throws IOException {

    SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions());

    return ReadPayload.parseFrom(

        PTransformTranslation.toProto(transform, Collections.emptyList(), components)

            .getSpec()

            .getPayload());

  }



  private static SdkFunctionSpec toProto(UnboundedSource<?, ?> source, SdkComponents components) {

    return SdkFunctionSpec.newBuilder()

        .setEnvironmentId(components.getOnlyEnvironmentId())

        .setSpec(

            FunctionSpec.newBuilder()

                .setUrn(JAVA_SERIALIZED_UNBOUNDED_SOURCE)

                .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(source)))

                .build())

        .build();

  }



  public static UnboundedSource<?, ?> unboundedSourceFromProto(ReadPayload payload)

      throws InvalidProtocolBufferException {

    checkArgument(payload.getIsBounded().equals(IsBounded.Enum.UNBOUNDED));

    return (UnboundedSource<?, ?>)

        SerializableUtils.deserializeFromByteArray(

            payload.getSource().getSpec().getPayload().toByteArray(), "UnboundedSource");

  }



  public static PCollection.IsBounded sourceIsBounded(AppliedPTransform<?, ?, ?> transform) {

    try {

      SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions());

      return PCollectionTranslation.fromProto(

          ReadPayload.parseFrom(

                  PTransformTranslation.toProto(transform, Collections.emptyList(), components)

                      .getSpec()

                      .getPayload())

              .getIsBounded());

    } catch (IOException e) {

      throw new RuntimeException("Internal error determining boundedness of Read", e);

    }

  }



  /** A {@link PTransformTranslation.TransformPayloadTranslator} for {@link Read.Unbounded}. */

  public static class UnboundedReadPayloadTranslator

      implements PTransformTranslation.TransformPayloadTranslator<Read.Unbounded<?>> {

    public static PTransformTranslation.TransformPayloadTranslator create() {

      return new UnboundedReadPayloadTranslator();

    }



    private UnboundedReadPayloadTranslator() {}



    @Override

    public String getUrn(Read.Unbounded<?> transform) {

      return PTransformTranslation.READ_TRANSFORM_URN;

    }



    @Override

    public FunctionSpec translate(

        AppliedPTransform<?, ?, Read.Unbounded<?>> transform, SdkComponents components) {

      ReadPayload payload = toProto(transform.getTransform(), components);

      return RunnerApi.FunctionSpec.newBuilder()

          .setUrn(getUrn(transform.getTransform()))

          .setPayload(payload.toByteString())

          .build();

    }

  }



  /** A {@link PTransformTranslation.TransformPayloadTranslator} for {@link Read.Bounded}. */

  public static class BoundedReadPayloadTranslator

      implements PTransformTranslation.TransformPayloadTranslator<Read.Bounded<?>> {

    public static PTransformTranslation.TransformPayloadTranslator create() {

      return new BoundedReadPayloadTranslator();

    }



    private BoundedReadPayloadTranslator() {}



    @Override

    public String getUrn(Read.Bounded<?> transform) {

      return PTransformTranslation.READ_TRANSFORM_URN;

    }



    @Override

    public FunctionSpec translate(

        AppliedPTransform<?, ?, Read.Bounded<?>> transform, SdkComponents components) {

      ReadPayload payload = toProto(transform.getTransform(), components);

      return RunnerApi.FunctionSpec.newBuilder()

          .setUrn(getUrn(transform.getTransform()))

          .setPayload(payload.toByteString())

          .build();

    }

  }



  /** Registers {@link UnboundedReadPayloadTranslator} and {@link BoundedReadPayloadTranslator}. */

  @AutoService(TransformPayloadTranslatorRegistrar.class)

  public static class Registrar implements TransformPayloadTranslatorRegistrar {

    @Override

    public Map<? extends Class<? extends PTransform>, ? extends PTransformTranslation.TransformPayloadTranslator>

        getTransformPayloadTranslators() {

      return ImmutableMap.<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder()

          .put(Read.Unbounded.class, new UnboundedReadPayloadTranslator())

          .put(Read.Bounded.class, new BoundedReadPayloadTranslator())

          .build();

    }

  }

}